from pinhole.models.base import EmbeddingModel

from dataclasses import dataclass, field
from os import environ
from typing import List

import requests


@dataclass
class BaichuanEmbeddingModel(EmbeddingModel):

    api_key: str = environ.get("BAICHUAN_API_KEY", "")

    def embed(self, input: str) -> List[float]:
        response = requests.post(
            "http://api.baichuan-ai.com/v1/embeddings",
            headers={
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            },
            json={
                "input": input,
                "model": "Baichuan-Text-Embedding"
            }
        )

        response_json = response.json()
        return response_json['data'][0]['embedding']
